//
// Advanced Java Programming
// TIEI - Fall 2024
//

/**
 * @author liangqiaowei & luzeyu
 */
package tiei.ajp.huffman;

import org.junit.Test;
import tiei.ajp.huffman.utils.*;

import java.io.*;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * A class to compress a file
 */
public class CompressFile {

    File inputFile = null;
    File outputFile = null;

    public CompressFile(File inFile, File outFile) {
        this.inputFile = inFile;
        this.outputFile = outFile;
    }

    private static final int EOF = -1;

    /**
     * Return an CompressFile object to compress 'inFile' into 'outFile'
     *
     * @param inFile  Input source file to be compressed
     * @param outFile Output destination file
     * @throws IOException
     */
    public static void compressFile(File inFile, File outFile) throws IOException {
//        assert inFile.canRead() && outFile.canWrite();
        BufferedInputStream bfIn = new BufferedInputStream(new FileInputStream(inFile));
        // ----- Cal freq num of each byte, to be used for encoding
        Map<Byte, Integer> byteFreqMap = new HashMap<>();
        while (true) {
            int data = bfIn.read();
            if (data == EOF) break;  // Break when reach EOF
            byte b = (byte) (data & BinaryUtils.LBYTE_MASK);
            if (byteFreqMap.containsKey(b)) {
                byteFreqMap.put(b, byteFreqMap.get(b) + 1);
            } else {
                byteFreqMap.put(b, 1);
            }
        }
        bfIn.close();
        // ----- Build huffman tree & generate encode map
        HuffmanTree<Byte> htree = new HuffmanTree<>();
        for (byte b : byteFreqMap.keySet()) {
            htree.addNode(new TreeNode<Byte>(byteFreqMap.get(b), b));
        }
        TreeNode<Byte> root = htree.buildTree();
        HuffmanEncoder<Byte> henc = new HuffmanEncoder<>();
        Map<Byte, String> encMap = henc.encode(root);
        // ----- Serialize huffman tree
        List<Byte> serializedTree = htree.serializeTree();
        // ----- Write all data
        BufferedOutputStream bfOut = new BufferedOutputStream(new FileOutputStream(outFile));
        // --- Write magic number
        bfOut.write(0x04);   // Magic num length
        bfOut.write(0x15);   // Magic num in little-endian, LSByte
        bfOut.write(0xcd);
        bfOut.write(0x5b);
        bfOut.write(0x07);   // Magic num in little-endian, HSByte
        // --- Write tree
        for (Byte b : serializedTree) {
            bfOut.write((int) b);
        }
        // --- Write bits num
        long bitNum = 0L;
        // Calc bit num using byteFreqMap and encMap
        for (Byte b : byteFreqMap.keySet()) {
            // byte freq * encoding len
            bitNum += (long) byteFreqMap.get(b) * (encMap.get(b).length());
        }
        // Calc `b`
        byte numB = (byte) (bitNum >= (1L << 56) ? 8 :
                bitNum >= (1L << 48) ? 7 :
                        bitNum >= (1L << 40) ? 6 :
                                bitNum >= (1L << 32) ? 5 :
                                        bitNum >= (1L << 24) ? 4 :
                                                bitNum >= (1L << 16) ? 3 :
                                                        bitNum >= (1L << 8) ? 2 :
                                                                bitNum >= (1L << 0) ? 1 : 0);
        assert numB != 0;
        // Write `b`
        bfOut.write((int) numB);
        // Write `bitNum` in little endian
        for (int i = 0; i < numB; i++) {
            byte lsb = (byte) ((bitNum >> (8 * i)) & 0xff);
            bfOut.write(lsb);
        }
        // --- Encode & serialize compressed data
        BitOutputStream bos = new BitOutputStream(bfOut);
        bfIn = new BufferedInputStream(new FileInputStream(inFile));
        while (true) {
            int data = bfIn.read();
            if (data == EOF) break;   // Break when EOF
            byte b = (byte) (data & BinaryUtils.LBYTE_MASK);
            String codeStr = encMap.get(b);
            for (int i = 0; i < codeStr.length(); i++) {
                char c = codeStr.charAt(i);
                assert c == '0' || c == '1';
                if (c == '0') {
                    bos.writeBit(BitOutputStream.BIT_O);
                } else {
                    bos.writeBit(BitOutputStream.BIT_1);
                }
            }
        }
        bos.flush();
        bfOut.flush();
        bfOut.close();
    }

    public void compress() throws IOException {
        compressFile(inputFile, outputFile);
    }

}
